from typing import List, Dict, Tuple
import numpy as np
from app.schemas.medical_share_schema import *
from app.tools.redis_client import redis_client
from app.utils.logger import logger

class MedicalShareService:
    def __init__(self):
        self.epsilon = 1e-10  # 避免除零错误的小值
        
    async def update_task_status(self, task_id: str, task_data: dict, progress: int, logs: str):
        """更新任务状态"""
        try:
            response = AlgorithmMiddleResponse(
                task_id=task_id,
                task_callback_url=task_data.get("task_callback_url"),
                task_status=TaskStatus.RUNNING,
                task_progress=progress,
                task_logs=logs,
                input_params=task_data.get("input_params"),
                metrics=[]
            )
            await redis_client.update_data(f'{task_id}_result', response.model_dump())
            logger.info(f"任务状态更新成功: {task_id}, 进度: {progress}%, 日志: {logs}")
        except Exception as e:
            logger.error(f"任务状态更新失败: {task_id}, 错误: {str(e)}")
            raise

    def calculate_privacy_score(self, strategy: SharingStrategy, party_info: DataPartyInfo) -> float:
        """计算隐私保护分数"""
        share_ratio = strategy.share_ratio
        privacy_budget = party_info.privacy_budget
        return 1 - (share_ratio / privacy_budget)

    def calculate_utility_score(self, strategies: List[SharingStrategy], parties: List[DataPartyInfo]) -> float:
        """计算数据效用分数"""
        total_shared = sum(s.share_ratio * p.data_size for s, p in zip(strategies, parties))
        total_possible = sum(p.data_size for p in parties)
        return total_shared / total_possible if total_possible > 0 else 0

    def calculate_fairness_score(self, strategies: List[SharingStrategy]) -> float:
        """计算公平性分数"""
        share_ratios = [s.share_ratio for s in strategies]
        if not share_ratios:
            return 0
        mean_ratio = np.mean(share_ratios)
        variance = np.mean([(r - mean_ratio) ** 2 for r in share_ratios])
        return 1 / (1 + variance)

    def optimize_sharing_strategy(
        self, 
        parties: List[DataPartyInfo],
        weights: Dict[ObjectiveType, float],
        current_strategies: List[SharingStrategy]
    ) -> Tuple[List[SharingStrategy], Dict[ObjectiveType, float]]:
        """优化数据共享策略"""
        new_strategies = []
        
        for i, party in enumerate(parties):
            # 计算新的共享比例
            current_ratio = current_strategies[i].share_ratio
            privacy_factor = weights[ObjectiveType.PRIVACY]
            utility_factor = weights[ObjectiveType.UTILITY]
            
            # 使用梯度下降更新共享比例
            gradient = utility_factor - privacy_factor
            new_ratio = current_ratio + 0.1 * gradient
            new_ratio = max(party.min_share_ratio, min(1.0, new_ratio))
            
            # 创建新的共享策略
            new_strategy = SharingStrategy(
                party_id=party.party_id,
                share_ratio=new_ratio,
                selected_features=current_strategies[i].selected_features,
                privacy_score=self.calculate_privacy_score(current_strategies[i], party)
            )
            new_strategies.append(new_strategy)
        
        # 计算目标函数值
        objective_scores = {
            ObjectiveType.PRIVACY: np.mean([self.calculate_privacy_score(s, p) 
                                          for s, p in zip(new_strategies, parties)]),
            ObjectiveType.UTILITY: self.calculate_utility_score(new_strategies, parties),
            ObjectiveType.FAIRNESS: self.calculate_fairness_score(new_strategies),
            ObjectiveType.EFFICIENCY: 1.0  # 简化处理
        }
        
        return new_strategies, objective_scores

    async def run_medical_share(self, input_params: dict, task_id: str) -> AlgorithmResponse:
        """运行医疗数据共享优化算法的主函数"""
        task_data = await redis_client.get_data(task_id)
        
        try:
            # 添加详细的日志来检查输入参数
            logger.info(f"收到的原始输入参数: {input_params}")
            
            # 检查medical_share_params是否存在
            medical_share_params = input_params.get("medical_share_params")
            if not medical_share_params:
                logger.error("缺少medical_share_params")
                raise ValueError("缺少medical_share_params")
            
            logger.info(f"解析的medical_share_params: {medical_share_params}")
            
            # 使用pydantic模型验证参数
            try:
                params = MedicalShareInputParams(**medical_share_params)
                logger.info(f"参数验证成功: {params}")
            except Exception as e:
                logger.error(f"参数验证失败: {str(e)}")
                raise
            
            await self.update_task_status(task_id, task_data, 10, "开始运行医疗数据共享优化算法...")
            
            # 初始化策略
            current_strategies = [
                SharingStrategy(
                    party_id=party.party_id,
                    share_ratio=party.min_share_ratio,
                    selected_features=[],
                    privacy_score=1.0
                ) for party in params.data_parties
            ]
            
            results = []
            convergence_history = []
            
            # 迭代优化
            for iteration in range(params.max_iterations):
                new_strategies, objective_scores = self.optimize_sharing_strategy(
                    params.data_parties,
                    params.objective_weights,
                    current_strategies
                )
                
                # 计算收敛度
                convergence = np.mean([abs(new.share_ratio - old.share_ratio) 
                                     for new, old in zip(new_strategies, current_strategies)])
                convergence_history.append(convergence)
                
                # 记录当前迭代结果
                results.append(TimeStepResult(
                    iteration=iteration,
                    strategies=new_strategies,
                    objective_scores=objective_scores,
                    pareto_optimal=convergence < params.convergence_threshold
                ))
                
                if convergence < params.convergence_threshold:
                    break
                    
                current_strategies = new_strategies
                await self.update_task_status(
                    task_id, 
                    task_data, 
                    min(90, int(10 + 80 * iteration / params.max_iterations)),
                    f"完成第{iteration + 1}次迭代，收敛度: {convergence:.6f}"
                )
            
            # 构建输出结果
            output_params = MedicalShareOutputParams(
                results=results,
                final_strategies=current_strategies,
                convergence_history=convergence_history,
                parameters={
                    "max_iterations": params.max_iterations,
                    "convergence_threshold": params.convergence_threshold,
                    "privacy_threshold": params.privacy_threshold,
                    "objective_weights": params.objective_weights
                }
            )
            
            return AlgorithmResponse(
                task_id=task_id,
                task_callback_url=task_data.get("task_callback_url"),
                task_status=TaskStatus.COMPLETED,
                task_progress=100,
                output_params=OutputParams(medical_share_results=output_params),
                metrics=[]
            )
            
        except Exception as e:
            logger.error(f"医疗数据共享优化算法执行失败: {str(e)}")
            error_response = AlgorithmMiddleResponse(
                task_id=task_id,
                task_callback_url=task_data.get("task_callback_url"),
                task_status=TaskStatus.FAILED,
                task_progress=0,
                task_logs=f"错误: {str(e)}",
                input_params=task_data.get("input_params"),
                metrics=[]
            )
            await redis_client.update_data(f'{task_id}_result', error_response.model_dump())
            raise 